import matplotlib.pyplot as plt
import torch
import SettingUp as s
import BuildingTheNetWorking as B

examples = enumerate(s.test_loader)
batch_idx, (example_data, example_targets) = next(examples)
with torch.no_grad():
  output = B.network(example_data)
fig = plt.figure()
for i in range(6):
  plt.subplot(2,3,i+1)
  plt.tight_layout()
  plt.imshow(example_data[i][0], cmap='gray', interpolation='none')
  plt.title("Prediction: {}".format(
    output.data.max(1, keepdim=True)[1][i].item()))
  plt.xticks([])
  plt.yticks([])
plt.show()